# -*- coding: utf-8 -*-  
'''
nlu相关评价

Created on 2021年10月2日
@author: luoyi
'''
import tensorflow as tf

import utils.logger_factory as logf
from utils.relationship import Relationship
from models.crf.process import padding_mask, viterbi


#    crf评价函数
class CRFMetric(tf.metrics.Metric):
    def __init__(self,
                 name='crf_metric',
                 crf_layer=None,
                 **kwargs):
        super(CRFMetric, self).__init__(name=name, **kwargs)
        
        self._crf_layer = crf_layer
        
        self._acc = self.add_weight('crf_metrics', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''计算词性标注预测正确的数量
            TP = et_to_et预测正确的数量
            TF = et_no预测正确的数量
            acc = TP + TF / P
        '''
        tf.print('==================================================================', output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('------------------------------ crf ------------------------------', output_stream=logf.get_logger_filepath('nlu_metrics'))
        
        #    取真实基序列
        y_id_pos = y_true[:, 1, :]
        
        #    取标注预测
        x_pos = self._crf_layer.get_out()
        #    状态转移矩阵
        transfer = self._crf_layer.get_transfer()
        
        pad_mask = padding_mask(y_id_pos)
        #    vierbi算法返回的路径不包含首尾
        x_id_pos, _, _, _ = viterbi(x_pos, pad_mask, transfer)
        
        #    每个batch下标注数（包含[CLS][EOS]）
        P_B = tf.math.count_nonzero(y_id_pos >= 0, axis=-1)
        y_id_pos = tf.RaggedTensor.from_tensor(y_id_pos, padding=-1)
        #    y_id_pos相应的也去掉首尾[CLS][POS]
        y_id_pos = y_id_pos[:, 1:-1]
        
        eq = tf.math.equal(x_id_pos, y_id_pos)
        eq = tf.RaggedTensor.to_tensor(eq, default_value=False)
        TP_B = tf.math.count_nonzero(eq, axis=-1)
        
        #    计算准确率
        acc = tf.math.reduce_sum(TP_B) / tf.math.reduce_sum(P_B)
        self._acc.assign(acc)
        
        tf.print('P_B', P_B, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('TP_B', TP_B, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('acc', acc, output_stream=logf.get_logger_filepath('nlu_metrics'))
        
        pass
    
    def reset_states(self):
        self._acc.assign(0.)
        pass
    def result(self):
        return self._acc
    pass


#    意图识别评价函数
class InformMetric(tf.keras.metrics.Metric):
    def __init__(self,
                 name='inform_f1',
                 inform_layer=None,
                 inform_size=Relationship.instance().rel_size(),
                 **kwargs):
        super(InformMetric, self).__init__(name=name, **kwargs)
        
        self._inform_layer = inform_layer
        self._inform_size = inform_size
        
        self._f1 = self.add_weight('inform_f1', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''计算意图识别的f1得分
            precision = 真实标签预测正确的数量
            recall = 预测标签预测正确的数量
            f1 = 2 * (precision * recall) / (precision + recall)
            
            @param y_true: Tensor(batch_size, 2, match_size)
        '''
        tf.print('------------------------------ f1_score ------------------------------', output_stream=logf.get_logger_filepath('nlu_metrics'))
        self.f1_score(y_true)
        pass
    
    #    计算预测准确率
    def account(self, y_true):
        #    取数
        y_true = y_true[:, 0, 0]                        #    Tensor(batch_size, )
        y_pred = self._inform_layer.get_out()           #    Tebsor(batch_size, inform_size)
        y_pred = tf.math.argmax(y_pred, axis=-1)        #    Tebsor(batch_size, )
        
        T = tf.math.count_nonzero(tf.math.equal(y_true, y_pred))
        N = tf.math.count_nonzero(y_true >= 0)
        acc = T / N
        self._f1.assign(acc)
        pass
    
    #    计算f1得分
    def f1_score(self, y_true):
        #    取数
        y_true = y_true[:, 0, 0]                        #    Tensor(batch_size, )
        y_pred = self._inform_layer.get_out()           #    Tebsor(batch_size, inform_size)
        y_pred = tf.math.argmax(y_pred, axis=-1)        #    Tebsor(batch_size, )
        
        #    本轮batch中每个意图的真实数量    Tensor(inform_size, )
        Ntrue = self.count(y_true, depty=self._inform_size)
        Ntrue = tf.cast(Ntrue, dtype=tf.float64)
        #    本轮batch中每个意图的预测数量    Tensor(inform_size, )
        Npred = self.count(y_pred, depty=self._inform_size) 
        Npred = tf.cast(Npred, dtype=tf.float64)
        #    本轮batch中每个意图预测正确的数量    Tensor(inform_size, )
        Tpred = tf.where(tf.math.equal(y_pred, y_true), y_true, -tf.ones_like(y_true))
        Tpred = Tpred[Tpred >= 0]
        Tpred = self.count(Tpred, depty=self._inform_size)
        Tpred = tf.cast(Tpred, dtype=tf.float64)
        
        #    计算精确率，-1表示本轮batch中y_true不含该意图
        precision = tf.where(Ntrue > 0, Tpred / Ntrue, tf.zeros_like(Ntrue))
        #    计算召回率，-1表示本轮batch中y_pred不含该意图
        recall = tf.where(Npred > 0, Tpred / Npred, tf.zeros_like(Npred))
        #    计算f1得分
        sum_precision_recall = precision + recall
        f1 = 2 * (precision * recall) / sum_precision_recall
        f1 = tf.where(tf.math.is_inf(f1), tf.zeros_like(f1), f1)
        f1 = tf.where(tf.math.is_nan(f1), tf.zeros_like(f1), f1)
        #    计算平均f1
        f1_mean = tf.math.reduce_mean(f1)
        
        tf.print('Ntrue:', Ntrue, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('Npred:', Npred, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('Tpred:', Tpred, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('precision:', precision, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('recall:', recall, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('f1:', f1, output_stream=logf.get_logger_filepath('nlu_metrics'))
        tf.print('f1_mean:', f1_mean, output_stream=logf.get_logger_filepath('nlu_metrics'))
        
        self._f1.assign(f1_mean)
        pass
    
    #    计算每个idx对应个数
    def count(self, vec, depty=Relationship.instance().rel_size()):
        #    计算每个元素个数
        uy, _,  uc = tf.unique_with_counts(vec)
        #    按索引排序，SparseTensor需要索引是有序的
        sort_idx = tf.argsort(uy, axis=-1)
        uy = tf.gather(uy, indices=sort_idx, axis=-1)
        uc = tf.gather(uc, indices=sort_idx, axis=-1)
        idx = tf.stack([tf.zeros_like(uy), uy], axis=-1)
        #    组成稀疏阵
        t = tf.SparseTensor(indices=tf.cast(idx, tf.int64), values=uc, dense_shape=(1, self._inform_size))
        t = tf.sparse.to_dense(t, default_value=0)[0]
        return t
    
    def reset_states(self):
        self._f1.assign(0.)
        pass
    def result(self):
        return self._f1
    pass



